package service

import (
	"fmt"

	"github.com/tetratelabs/wazero/api"
	"github.com/tokentransfer/node/core"
	"github.com/tokentransfer/node/util"
)

const (
	MALLOC_FUNCTION = "allocate"
	FREE_FUNCTION   = "deallocate"

	START_FUNCTION   = "start"
	PROCESS_FUNCTION = "process"
	STOP_FUNCTION    = "stop"
)

func verifyTypes(name string, requiredTypes []api.ValueType, definedTypes []api.ValueType) error {
	rl := len(requiredTypes)
	dl := len(definedTypes)
	if rl != dl {
		return util.ErrorOfUnmatched("length of types", name, rl, dl)
	}
	for i := 0; i < rl; i++ {
		rt := requiredTypes[i]
		dt := definedTypes[i]
		if rt != dt {
			return util.ErrorOfUnmatched(fmt.Sprintf("%d type", i), name, rt, dt)
		}
	}
	return nil
}

func verifyFunc(mod api.Module, name string, inputs []api.ValueType, outputs []api.ValueType) (api.Function, error) {
	f := mod.ExportedFunction(name)
	if f == nil {
		return nil, util.ErrorOfNotFound(name, "function")
	}
	err := verifyTypes(fmt.Sprintf("inputs of %s", name), inputs, f.Definition().ParamTypes())
	if err != nil {
		return nil, err
	}
	err = verifyTypes(fmt.Sprintf("outputs of %s", name), outputs, f.Definition().ResultTypes())
	if err != nil {
		return nil, err
	}
	return f, nil
}

func VerifyCode(s core.Storage, wasmCode []byte) (*WasmModule, error) {
	cost := api.NewCost()
	wm := &WasmModule{s: s}

	mod, err := wm.Load(cost, wasmCode)
	if err != nil {
		return nil, err
	}
	mallocFunc, err := verifyFunc(mod, MALLOC_FUNCTION, []api.ValueType{api.ValueTypeI32}, []api.ValueType{api.ValueTypeI64})
	if err != nil {
		return nil, err
	}
	freeFunc, err := verifyFunc(mod, FREE_FUNCTION, []api.ValueType{api.ValueTypeI64, api.ValueTypeI32}, []api.ValueType{})
	if err != nil {
		return nil, err
	}
	startFunc, err := verifyFunc(mod, START_FUNCTION, []api.ValueType{}, []api.ValueType{})
	if err != nil {
		return nil, err
	}
	stopFunc, err := verifyFunc(mod, STOP_FUNCTION, []api.ValueType{}, []api.ValueType{})
	if err != nil {
		return nil, err
	}
	processFunc, err := verifyFunc(mod, PROCESS_FUNCTION, []api.ValueType{api.ValueTypeI64}, []api.ValueType{api.ValueTypeI64})
	if err != nil {
		return nil, err
	}

	wm.mod = mod
	wm.mallocFn = mallocFunc
	wm.freeFn = freeFunc
	wm.startFn = startFunc
	wm.processFn = processFunc
	wm.stopFn = stopFunc

	return wm, nil
}
